# -*- coding: utf-8 -*-

import os
import cv2
import oyaml as yaml


color_palette = [(255, 0, 0), (0, 255, 0), (0, 0, 255),
                 (255, 255, 0), (255, 0, 255), (0, 255, 255)]


class LabelVideoBackend(object):
    def __init__(self, root_folder, video_path):
        self._root_folder = root_folder
        self._video_path = video_path
        self._video = cv2.VideoCapture(video_path)
        self._door_open_frame_information = self.generate_door_open_frame_information()
        self._door_open_frame_indices = list(self._door_open_frame_information.keys())
        self._door_open_frame_indices_index = 0

    def get_door_open_frame_indices(self):  # 人工标注开关门的结果
        yaml_file_path = os.path.join(self._root_folder, "information", "label", "door.yaml")
        with open(yaml_file_path, encoding="utf-8") as yaml_file:
            data = yaml.load(yaml_file)
            video_door_status = data["video_door_status"]
        video_path = self._video_path.replace(self._root_folder + "\\", "")
        door_open_frame_segments = video_door_status[video_path]["开门"]

        frame_indices = list()
        for segment in door_open_frame_segments:
            start, end = segment
            if end > self._video.get(cv2.CAP_PROP_FRAME_COUNT):
                end = self._video.get(cv2.CAP_PROP_FRAME_COUNT) - 1
            frame_indices.extend(list(range(start, end+1)))
        return frame_indices

    def get_labeled_information(self):  # 算法预测和人工标注的结果
        yaml_file_path = os.path.join(self._root_folder, "information", "label", "flow_label.yaml")
        with open(yaml_file_path, encoding="utf-8") as yaml_file:
            data = yaml.load(yaml_file)
            video_labeled_positions = data["labeled_positions"]
        video_path = self._video_path.replace(self._root_folder + "\\", "")
        if video_path in video_labeled_positions:
            labeled_information = video_labeled_positions[video_path]
            return labeled_information
        return None

    def generate_door_open_frame_information(self):
        door_open_frame_indices = self.get_door_open_frame_indices()
        labeled_information = self.get_labeled_information()
        door_open_frame_information = {frame_index: {} for frame_index in door_open_frame_indices}
        if labeled_information is not None:
            for frame_index in door_open_frame_information:
                if frame_index in labeled_information:
                    door_open_frame_information[frame_index] = labeled_information[frame_index]
        return door_open_frame_information

    @staticmethod
    def draw_rectangle(frame, person_id, position):
        color = color_palette[person_id % len(color_palette)]
        x1, y1, x2, y2 = position
        x1 = int(x1 * 320)
        y1 = int(y1 * 288)
        x2 = int(x2 * 320)
        y2 = int(y2 * 288)
        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
        cv2.putText(frame, str(person_id), (x1+10, y1+20), 1, 1.5, color)

    def load_image(self):
        frame_index = self._door_open_frame_indices[self._door_open_frame_indices_index]
        self._video.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
        res, frame = self._video.read()
        labeled_information = self._door_open_frame_information[frame_index]
        if len(labeled_information) > 0:
            for person_id, position in labeled_information.items():
                self.draw_rectangle(frame, person_id, position)
        return frame

    def remove_one_rectangle(self):
        frame_index = self._door_open_frame_indices[self._door_open_frame_indices_index]
        frame_information = self._door_open_frame_information[frame_index]
        if len(frame_information) > 0:
            key = list(frame_information.keys())[-1]
            frame_information.pop(key)

    def add_one_rectangle(self, rectangle):
        frame_index = self._door_open_frame_indices[self._door_open_frame_indices_index]
        frame_information = self._door_open_frame_information[frame_index]
        if len(frame_information) > 0:
            person_id = list(frame_information.keys())[-1] + 1
        else:
            person_id = 1
        frame_information[person_id] = rectangle

    def save_labeled_result(self):
        yaml_file_path = os.path.join(self._root_folder, "information", "label", "flow_label.yaml")
        with open(yaml_file_path, encoding="utf-8") as yaml_file:
            data = yaml.load(yaml_file)
        video_path = self._video_path.replace(self._root_folder + "\\", "")
        if video_path in data["labeled_positions"]:
            data["labeled_positions"][video_path] = dict()
            for frame_index, person_positions in self._door_open_frame_information.items():
                if len(person_positions) > 0:
                    data["labeled_positions"][video_path][frame_index] = person_positions
        with open(yaml_file_path, "w", encoding="utf-8") as yaml_file:
            yaml.dump(
                data,
                yaml_file,
                allow_unicode=True,
                default_flow_style=False)

    @property
    def labeled_frame_number(self):
        return len(self._door_open_frame_indices)

    @property
    def labeled_frame_index(self):
        return self._door_open_frame_indices_index

    @labeled_frame_index.setter
    def labeled_frame_index(self, index):
        if 0 <= index < len(self._door_open_frame_indices):
            self._door_open_frame_indices_index = index


def main():
    root_folder = r"F:\data\bus_videos\59092_00E0B453CE0B\2019_3_18"
    video_path = r"F:\data\bus_videos\59092_00E0B453CE0B\2019_3_18\c0\2019_3_18-7_5_7-0(9_0_9_0).avi"
    backend = LabelVideoBackend(root_folder, video_path)


if __name__ == "__main__":
    main()
